import asyncio
import logging
from typing import Dict, Literal, Optional

import orjson
from cachetools import TTLCache
from langfuse.decorators import observe
from pydantic import BaseModel

from src.core.pipeline import BasicPipeline
from src.utils import trace_metadata
from src.web.v1.services import Configuration, MetadataTraceable

logger = logging.getLogger("wren-ai-service")


class QuestionRecommendation:
    class Input(BaseModel):
        id: str
        mdl: str
        previous_questions: list[str] = []
        project_id: Optional[str] = None
        max_questions: Optional[int] = 5
        max_categories: Optional[int] = 3
        regenerate: Optional[bool] = False
        configuration: Optional[Configuration] = Configuration()

    class Resource(BaseModel, MetadataTraceable):
        class Error(BaseModel):
            code: Literal["OTHERS", "MDL_PARSE_ERROR", "RESOURCE_NOT_FOUND"]
            message: str

        id: str
        status: Literal["generating", "finished", "failed"] = "generating"
        response: Optional[dict] = {"questions": {}}
        error: Optional[Error] = None
        trace_id: Optional[str] = None

    def __init__(
        self,
        pipelines: Dict[str, BasicPipeline],
        maxsize: int = 1_000_000,
        ttl: int = 120,
    ):
        self._pipelines = pipelines
        self._cache: Dict[str, QuestionRecommendation.Resource] = TTLCache(
            maxsize=maxsize, ttl=ttl
        )

    def _handle_exception(
        self,
        input: Input,
        error_message: str,
        code: str = "OTHERS",
        trace_id: Optional[str] = None,
    ):
        self._cache[input.id] = self.Resource(
            id=input.id,
            status="failed",
            error=self.Resource.Error(code=code, message=error_message),
            trace_id=trace_id,
        )
        logger.error(error_message)

    @observe(name="Validate Question")
    async def _validate_question(
        self,
        candidate: dict,
        request_id: str,
        max_questions: int,
        max_categories: int,
        project_id: Optional[str] = None,
        configuration: Optional[Configuration] = Configuration(),
    ):
        try:
            # 检查candidate是否有效
            if not isinstance(candidate, dict):
                logger.error(f"Invalid candidate format: {candidate}")
                return {}
                
            # 确保candidate有question字段
            if "question" not in candidate:
                logger.error(f"Candidate missing 'question' field: {candidate}")
                return {}
                
            # 确保candidate有category字段
            if "category" not in candidate:
                logger.info(f"Candidate missing 'category' field, using 'General': {candidate}")
                candidate["category"] = "General"
                
            logger.info(f"Validating question: {candidate['question']} (Category: {candidate['category']})")
            
            # 执行检索
            try:
                retrieval_result = await self._pipelines["retrieval"].run(
                    query=candidate["question"],
                    id=project_id,
                )
                
                if not retrieval_result:
                    logger.error(f"Retrieval returned None for question: {candidate['question']}")
                    return {}
                    
                _retrieval_result = retrieval_result.get("construct_retrieval_results", {})
                documents = _retrieval_result.get("retrieval_results", [])
                table_ddls = [document.get("table_ddl") for document in documents]
                has_calculated_field = _retrieval_result.get("has_calculated_field", False)
                has_metric = _retrieval_result.get("has_metric", False)
                
                logger.info(f"Retrieved {len(documents)} documents for question: {candidate['question']}")
            except Exception as e:
                logger.error(f"Error during retrieval for question '{candidate['question']}': {str(e)}")
                return {}

            # 执行SQL生成推理
            try:
                sql_generation_reasoning = (
                    await self._pipelines["sql_generation_reasoning"].run(
                        query=candidate["question"],
                        contexts=table_ddls,
                        configuration=configuration,
                    )
                ).get("post_process", {})
                
                logger.info(f"Generated SQL reasoning for question: {candidate['question']}")
            except Exception as e:
                logger.error(f"Error during SQL generation reasoning for question '{candidate['question']}': {str(e)}")
                return {}

            # 执行SQL生成
            try:
                generated_sql = await self._pipelines["sql_generation"].run(
                    query=candidate["question"],
                    contexts=table_ddls,
                    sql_generation_reasoning=sql_generation_reasoning,
                    configuration=configuration,
                    project_id=project_id,
                    has_calculated_field=has_calculated_field,
                    has_metric=has_metric,
                )
                
                if not generated_sql or not isinstance(generated_sql, dict):
                    logger.error(f"Invalid SQL generation result: {generated_sql}")
                    return {}
                    
                post_process = generated_sql.get("post_process", {})
                
                if not post_process:
                    logger.error(f"Empty post_process in SQL generation result")
                    return {}
                    
                logger.info(f"Generated SQL for question: {candidate['question']}")
            except Exception as e:
                logger.error(f"Error during SQL generation for question '{candidate['question']}': {str(e)}")
                return {}

            # 检查是否有有效的生成结果
            valid_results = post_process.get("valid_generation_results", [])
            if not valid_results:
                logger.warning(f"No valid SQL generated for question: {candidate['question']}")
                return post_process

            # 获取第一个有效SQL
            valid_sql = ""
            if isinstance(valid_results[0], dict):
                valid_sql = valid_results[0].get("sql", "")
            else:
                try:
                    # 尝试转换为字符串
                    valid_sql = str(valid_results[0])
                    logger.warning(f"Converted non-dict SQL result to string: {valid_sql}")
                except Exception as e:
                    logger.error(f"Failed to convert SQL result to string: {str(e)}")
                    return post_process
                
            if not valid_sql:
                logger.warning(f"Empty SQL generated for question: {candidate['question']}")
                return post_process

            logger.info(f"Valid SQL generated for question: {candidate['question']}")

            # 更新资源
            try:
                # 确保资源存在
                if request_id not in self._cache:
                    logger.error(f"Resource {request_id} not found")
                    return post_process
                    
                current = self._cache[request_id]
                
                # 确保response存在
                if current.response is None:
                    current.response = {"questions": {}}
                    
                # 确保questions存在
                if "questions" not in current.response:
                    current.response["questions"] = {}
                    
                questions = current.response["questions"]

                # 检查类别数量限制
                if (
                    candidate["category"] not in questions
                    and len(questions) >= max_categories
                ):
                    logger.info(f"Skipping question due to max categories limit: {candidate['question']}")
                    return post_process

                # 获取或创建类别
                current_category = questions.setdefault(candidate["category"], [])

                # 检查问题数量限制
                if len(current_category) >= max_questions:
                    logger.info(f"Skipping question due to max questions limit for category '{candidate['category']}'")
                    return post_process

                # 添加问题并确保SQL是字符串
                try:
                    question_with_sql = {**candidate, "sql": valid_sql}
                    # 执行最终检查，确保SQL是字符串
                    if not isinstance(question_with_sql["sql"], str):
                        question_with_sql["sql"] = str(question_with_sql["sql"])
                    current_category.append(question_with_sql)
                    logger.info(f"Added question to category '{candidate['category']}': {candidate['question']}")
                except Exception as e:
                    logger.error(f"Error adding question: {str(e)}")
                    # 尝试添加不带SQL的问题
                    question_without_sql = {**candidate}
                    question_without_sql.pop("sql", None)
                    current_category.append(question_without_sql)
                    logger.warning(f"Added question without SQL due to error: {candidate['question']}")
                
                return post_process
            except Exception as e:
                logger.error(f"Error updating resource for question '{candidate['question']}': {str(e)}")
                return post_process

        except Exception as e:
            logger.error(f"Error validating question: {str(e)}")
            return {}

    async def _recommend(self, request: dict, input: Input):
        try:
            # 确保pipeline存在
            if "question_recommendation" not in self._pipelines:
                logger.error(f"Pipeline 'question_recommendation' not found")
                self._cache[input.id].status = "failed"
                self._cache[input.id].error = self.Resource.Error(
                    code="OTHERS", 
                    message="Pipeline 'question_recommendation' not found"
                )
                return
                
            # 运行pipeline
            resp = await self._pipelines["question_recommendation"].run(**request)
            
            # 检查resp是否为None
            if resp is None:
                logger.error(f"Pipeline returned None for request {input.id}")
                self._cache[input.id].status = "failed"
                self._cache[input.id].error = self.Resource.Error(
                    code="OTHERS", 
                    message="Pipeline returned None"
                )
                return
                
            # 检查resp的类型
            logger.info(f"Pipeline response type: {type(resp)}")
            
            # 安全记录内容，避免过大的日志
            try:
                if isinstance(resp, dict):
                    # 记录字典的键
                    logger.info(f"Pipeline response keys: {list(resp.keys())}")
                    # 尝试安全记录一部分内容
                    if "questions" in resp and isinstance(resp["questions"], list):
                        logger.info(f"Questions count: {len(resp['questions'])}")
                        if len(resp["questions"]) > 0:
                            sample = resp["questions"][0]
                            logger.info(f"First question sample: {sample}")
                elif isinstance(resp, list):
                    logger.info(f"Pipeline response list length: {len(resp)}")
                    if len(resp) > 0:
                        logger.info(f"First item sample: {resp[0]}")
                else:
                    # 对于其他类型，尝试获取字符串表示
                    logger.info(f"Pipeline response string rep: {str(resp)[:200]}...")
            except Exception as e:
                logger.warning(f"Error logging response content: {str(e)}")
            
            # 处理不同类型的响应
            questions = []
            if isinstance(resp, dict):
                # 如果是字典，尝试获取questions
                if "normalized" in resp and isinstance(resp["normalized"], dict):
                    questions = resp["normalized"].get("questions", [])
                elif "questions" in resp:
                    questions = resp["questions"]
            elif isinstance(resp, list):
                # 如果是列表，直接使用
                questions = resp
            else:
                # 其他类型，转换为字符串并尝试解析
                try:
                    str_resp = str(resp)
                    logger.info(f"Converting non-dict/list response to string: {str_resp[:100]}...")
                    # 尝试解析为JSON
                    try:
                        import json
                        json_resp = json.loads(str_resp)
                        if isinstance(json_resp, dict) and "questions" in json_resp:
                            questions = json_resp["questions"]
                        elif isinstance(json_resp, list):
                            questions = json_resp
                        else:
                            logger.warning(f"Could not extract questions from parsed JSON: {type(json_resp)}")
                    except json.JSONDecodeError:
                        logger.warning("Failed to parse response as JSON")
                except Exception as e:
                    logger.error(f"Error converting response to string: {str(e)}")
                    self._cache[input.id].status = "failed"
                    self._cache[input.id].error = self.Resource.Error(
                        code="OTHERS", 
                        message=f"Error processing response: {str(e)}"
                    )
                    return
                
            # 确保questions是列表
            if not isinstance(questions, list):
                logger.error(f"Questions is not a list: {type(questions)}")
                self._cache[input.id].status = "failed"
                self._cache[input.id].error = self.Resource.Error(
                    code="OTHERS", 
                    message=f"Questions is not a list: {type(questions)}"
                )
                return
                
            logger.info(f"Processing {len(questions)} questions for request {input.id}")
            
            # 如果questions为空，生成默认问题
            if len(questions) == 0:
                logger.warning(f"Empty questions list for request {input.id}, generating default questions")
                
                # 根据MDL生成默认问题
                mdl = request.get("mdl", {})
                models = mdl.get("models", [])
                language = request.get("language", "English")
                
                # 默认问题类别
                categories = ["Overview", "Trends", "Comparisons"]
                
                # 生成默认问题
                default_questions = []
                
                # 如果有模型，添加基于模型的问题
                if models and len(models) > 0:
                    for model in models:
                        model_name = model.get("name", "")
                        display_name = model.get("properties", {}).get("displayName", model_name)
                        columns = model.get("columns", [])
                        
                        if not model_name or not columns:
                            continue
                            
                        # 获取列名
                        column_names = [col.get("name") for col in columns if col.get("name")]
                        
                        # 根据语言生成问题
                        if language == "Simplified Chinese":
                            # 添加概览类问题
                            default_questions.append({
                                "question": f"显示{display_name}表中的所有数据",
                                "category": "Overview",
                                "sql": f"SELECT * FROM {model_name} LIMIT 10"
                            })
                            
                            # 如果有ID列，添加计数问题
                            if "id" in column_names:
                                default_questions.append({
                                    "question": f"{display_name}表中有多少条记录？",
                                    "category": "Overview",
                                    "sql": f"SELECT COUNT(*) FROM {model_name}"
                                })
                            
                            # 如果有文本列，添加分组问题
                            text_columns = [col for col in columns if col.get("type") in ["TEXT", "VARCHAR", "CHAR", "STRING"]]
                            if text_columns:
                                text_col = text_columns[0].get("name")
                                default_questions.append({
                                    "question": f"按{text_col}字段分组统计{display_name}表中的记录数",
                                    "category": "Comparisons",
                                    "sql": f"SELECT {text_col}, COUNT(*) FROM {model_name} GROUP BY {text_col}"
                                })
                        else:
                            # 添加概览类问题
                            default_questions.append({
                                "question": f"Show all data from {display_name}",
                                "category": "Overview",
                                "sql": f"SELECT * FROM {model_name} LIMIT 10"
                            })
                            
                            # 如果有ID列，添加计数问题
                            if "id" in column_names:
                                default_questions.append({
                                    "question": f"How many records are in the {display_name} table?",
                                    "category": "Overview",
                                    "sql": f"SELECT COUNT(*) FROM {model_name}"
                                })
                            
                            # 如果有文本列，添加分组问题
                            text_columns = [col for col in columns if col.get("type") in ["TEXT", "VARCHAR", "CHAR", "STRING"]]
                            if text_columns:
                                text_col = text_columns[0].get("name")
                                default_questions.append({
                                    "question": f"Count records in {display_name} grouped by {text_col}",
                                    "category": "Comparisons",
                                    "sql": f"SELECT {text_col}, COUNT(*) FROM {model_name} GROUP BY {text_col}"
                                })
                
                # 如果仍然没有问题，添加通用问题
                if not default_questions:
                    if language == "Simplified Chinese":
                        default_questions = [
                            {
                                "question": "显示所有表的数据",
                                "category": "Overview",
                                "sql": "SELECT * FROM information_schema.tables WHERE table_schema = 'public'"
                            },
                            {
                                "question": "数据库中有哪些表？",
                                "category": "Overview",
                                "sql": "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'"
                            },
                            {
                                "question": "显示数据库的基本信息",
                                "category": "Overview",
                                "sql": "SELECT current_database(), current_schema()"
                            }
                        ]
                    else:
                        default_questions = [
                            {
                                "question": "Show all tables in the database",
                                "category": "Overview",
                                "sql": "SELECT * FROM information_schema.tables WHERE table_schema = 'public'"
                            },
                            {
                                "question": "What tables are in the database?",
                                "category": "Overview",
                                "sql": "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'"
                            },
                            {
                                "question": "Show basic database information",
                                "category": "Overview",
                                "sql": "SELECT current_database(), current_schema()"
                            }
                        ]
                
                questions = default_questions
                logger.info(f"Generated {len(questions)} default questions for request {input.id}")
                
                # 直接更新资源的响应
                formatted_questions = {}
                for question in questions:
                    category = question.get("category", "General")
                    if category not in formatted_questions:
                        formatted_questions[category] = []
                    
                    # 确保SQL字段是字符串
                    if "sql" in question and not isinstance(question["sql"], str):
                        try:
                            # 如果SQL是dict或其他类型，尝试转换为字符串
                            if isinstance(question["sql"], dict):
                                question["sql"] = str(question["sql"].get("sql", ""))
                            else:
                                question["sql"] = str(question["sql"])
                        except Exception as e:
                            logger.warning(f"Error converting SQL to string: {e}")
                            question["sql"] = ""
                    
                    formatted_questions[category].append(question)
                
                self._cache[input.id].response = {"questions": formatted_questions}
                self._cache[input.id].status = "finished"
                
                logger.info(f"Updated response for request {input.id} with default questions")
                return

            # 验证每个问题
            validation_tasks = []
            for question in questions:
                # 确保question是字典
                if not isinstance(question, dict):
                    # 尝试转换非字典问题
                    try:
                        if isinstance(question, str):
                            # 尝试解析为JSON
                            try:
                                import json
                                parsed_question = json.loads(question)
                                if isinstance(parsed_question, dict):
                                    logger.info(f"Successfully parsed string question to dict: {parsed_question}")
                                    question = parsed_question
                                else:
                                    # 创建简单问题
                                    question = {"question": question, "category": "General"}
                                    logger.info(f"Created simple question from string: {question}")
                            except json.JSONDecodeError:
                                # 创建简单问题
                                question = {"question": question, "category": "General"}
                                logger.info(f"Created simple question from non-JSON string: {question}")
                        else:
                            # 尝试转换为字典
                            question_str = str(question)
                            question = {"question": question_str, "category": "General"}
                            logger.info(f"Created simple question from non-string: {question}")
                    except Exception as e:
                        logger.warning(f"Skipping non-dict question that couldn't be converted: {e}")
                        continue
                    
                # 确保question有必要的字段
                if "question" not in question:
                    logger.warning(f"Skipping question without 'question' field: {question}")
                    continue
                    
                validation_tasks.append(
                    self._validate_question(
                        question,
                        input.id,
                        input.max_questions,
                        input.max_categories,
                        input.project_id,
                        input.configuration,
                    )
                )

            if validation_tasks:
                await asyncio.gather(*validation_tasks, return_exceptions=True)
            else:
                logger.warning(f"No valid questions to validate for request {input.id}")
                self._cache[input.id].status = "finished"
            
        except Exception as e:
            logger.error(f"Error in _recommend for request {input.id}: {str(e)}")
            self._cache[input.id].status = "failed"
            self._cache[input.id].error = self.Resource.Error(
                code="OTHERS", 
                message=f"Error in _recommend: {str(e)}"
            )

    @observe(name="Generate Question Recommendation")
    @trace_metadata
    async def recommend(self, input: Input, **kwargs) -> Resource:
        logger.info(
            f"Request {input.id}: Generate Question Recommendation pipeline is running..."
        )
        trace_id = kwargs.get("trace_id")
        
        # 确保资源存在
        if input.id not in self._cache:
            self._cache[input.id] = self.Resource(id=input.id)

        try:
            # 解析MDL
            try:
                # 记录原始MDL以便调试
                logger.debug(f"Raw MDL for request {input.id}: {input.mdl[:100]}...")
                
                # 检查MDL是否为空
                if not input.mdl:
                    error_msg = "MDL is empty"
                    logger.error(error_msg)
                    self._cache[input.id].status = "failed"
                    self._cache[input.id].error = self.Resource.Error(
                        code="MDL_PARSE_ERROR", 
                        message=error_msg
                    )
                    self._cache[input.id].trace_id = trace_id
                    return self._cache[input.id].with_metadata()
                
                # 检查input.mdl是否已经是字典
                if isinstance(input.mdl, dict):
                    logger.info(f"MDL is already a dictionary for request {input.id}")
                    mdl_json = input.mdl
                else:
                    # 尝试解析MDL
                    try:
                        mdl_json = orjson.loads(input.mdl)
                    except orjson.JSONDecodeError as e:
                        # 如果解析失败，尝试修复常见问题
                        logger.warning(f"Initial JSON parse failed: {str(e)}. Trying to fix...")
                        
                        # 尝试修复单引号问题
                        fixed_mdl = input.mdl.replace("'", '"')
                        try:
                            mdl_json = orjson.loads(fixed_mdl)
                            logger.info("Successfully parsed MDL after fixing quotes")
                        except orjson.JSONDecodeError:
                            # 尝试移除空白字符
                            logger.warning("Quote fixing failed, trying to strip whitespace...")
                            try:
                                fixed_mdl = input.mdl.strip()
                                mdl_json = orjson.loads(fixed_mdl)
                                logger.info("Successfully parsed MDL after stripping whitespace")
                            except orjson.JSONDecodeError:
                                # 最后尝试使用ast.literal_eval
                                logger.warning("All JSON parsing methods failed, trying ast.literal_eval...")
                                import ast
                                try:
                                    mdl_json = ast.literal_eval(input.mdl)
                                    logger.info("Successfully parsed MDL using ast.literal_eval")
                                except (SyntaxError, ValueError) as ast_error:
                                    # 如果所有方法都失败，则抛出异常
                                    raise orjson.JSONDecodeError(f"Failed to parse MDL: {str(e)}", input.mdl, 0)
                
                logger.info(f"Successfully parsed MDL for request {input.id}")
            except orjson.JSONDecodeError as e:
                error_msg = f"Failed to parse MDL: {str(e)}"
                logger.error(error_msg)
                self._cache[input.id].status = "failed"
                self._cache[input.id].error = self.Resource.Error(
                    code="MDL_PARSE_ERROR", 
                    message=error_msg
                )
                self._cache[input.id].trace_id = trace_id
                return self._cache[input.id].with_metadata()
                
            # 准备请求
            request = {
                "mdl": mdl_json,
                "previous_questions": input.previous_questions or [],
                "language": input.configuration.language,
                "current_date": input.configuration.show_current_time(),
                "max_questions": input.max_questions,
                "max_categories": input.max_categories,
            }

            logger.info(f"Request data for {input.id}: {request}")

            # 第一阶段推荐
            await self._recommend(request, input)

            # 检查状态
            if self._cache[input.id].status == "failed":
                logger.error(f"First recommendation phase failed for request {input.id}")
                return self._cache[input.id].with_metadata()

            resource = self._cache[input.id]
            resource.trace_id = trace_id
            response = resource.response

            # 检查response是否有效
            if not response or not isinstance(response, dict) or "questions" not in response:
                logger.error(f"Invalid response format for request {input.id}: {response}")
                resource.status = "failed"
                resource.error = self.Resource.Error(
                    code="OTHERS", 
                    message="Invalid response format"
                )
                return resource.with_metadata()
                
            # 检查questions是否有效
            questions = response.get("questions", {})
            if not isinstance(questions, dict):
                logger.error(f"Invalid questions format for request {input.id}: {questions}")
                resource.status = "failed"
                resource.error = self.Resource.Error(
                    code="OTHERS", 
                    message="Invalid questions format"
                )
                return resource.with_metadata()
                
            # 如果没有问题，直接返回
            if not questions:
                logger.warning(f"No questions generated for request {input.id}")
                resource.status = "finished"
                return resource.with_metadata()

            # 计算需要补充的类别
            categories_count = {
                category: input.max_questions - len(category_questions)
                for category, category_questions in questions.items()
                if len(category_questions) < input.max_questions
            }
            categories = list(categories_count.keys())
            need_regenerate = len(categories) > 0 and input.regenerate

            resource.status = "generating" if need_regenerate else "finished"

            if resource.status == "finished":
                logger.info(f"Question recommendation completed for request {input.id}")
                return resource.with_metadata()

            # 第二阶段推荐
            try:
                await self._recommend(
                    {
                        **request,
                        "categories": categories,
                        "max_categories": len(categories),
                    },
                    input,
                )
            except Exception as e:
                logger.error(f"Second recommendation phase error for request {input.id}: {str(e)}")
                # 不将状态设为失败，因为第一阶段已经成功

            # 设置最终状态
            self._cache[input.id].status = "finished"
            logger.info(f"Question recommendation completed for request {input.id}")

        except orjson.JSONDecodeError as e:
            error_msg = f"Failed to parse MDL: {str(e)}"
            logger.error(error_msg)
            self._handle_exception(
                input,
                error_msg,
                code="MDL_PARSE_ERROR",
                trace_id=trace_id,
            )
        except Exception as e:
            error_msg = f"An error occurred during question recommendation generation: {str(e)}"
            logger.error(error_msg)
            self._handle_exception(
                input,
                error_msg,
                trace_id=trace_id,
            )

        return self._cache[input.id].with_metadata()

    def __getitem__(self, id: str) -> Resource:
        try:
            response = self._cache.get(id)

            if response is None:
                message = f"Question Recommendation Resource with ID '{id}' not found."
                logger.warning(message)
                new_resource = self.Resource(
                    id=id,
                    status="failed",
                    response={"questions": {}},
                    error=self.Resource.Error(
                        code="RESOURCE_NOT_FOUND",
                        message=message
                    )
                )
                self._cache[id] = new_resource
                return new_resource

            return response
        except Exception as e:
            error_message = f"Error retrieving resource with ID '{id}': {str(e)}"
            logger.error(error_message)
            new_resource = self.Resource(
                id=id,
                status="failed",
                response={"questions": {}},
                error=self.Resource.Error(
                    code="OTHERS",
                    message=error_message
                )
            )
            self._cache[id] = new_resource
            return new_resource

    def __setitem__(self, id: str, value: Resource):
        self._cache[id] = value